#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import math
from src.manager.log_manager import LogManager

Logger = LogManager.get_logger(__name__)

tf.compat.v1.disable_eager_execution()
tf.compat.v1.disable_v2_behavior()

class LstmManager:
    """
    长短期记忆网络管理器
    """

    def __init__(self):
        pass

    def test(self, data_frame):
        # importing required libraries
        from sklearn.preprocessing import MinMaxScaler
        from keras.models import Sequential
        from keras.layers import Dense, Dropout, LSTM

        # creating dataframe
        # data = data_frame.sort_index(ascending=True, axis=0)
        # new_data = pd.DataFrame(index=range(0, len(data_frame)), columns=['0', '1'])
        new_data = pd.DataFrame(data_frame)
        # for i in range(0, len(data)):
        #     new_data['0'][i] = data['0'][i]
        #     new_data['1'][i] = data['1'][i]


        # setting index
        # new_data.index = new_data.Date
        # new_data.drop('1', axis=1, inplace=True)

        # creating train and test sets
        dataset = new_data.values

        train = dataset[0:1500, :]
        valid = dataset[1500:, :]

        # converting dataset into x_train and y_train
        scaler = MinMaxScaler(feature_range=(0, 1))
        scaled_data = scaler.fit_transform(dataset)

        x_train, y_train = [], []
        for i in range(60, len(train)):
            x_train.append(scaled_data[i - 60:i, 0])
            y_train.append(scaled_data[i, 0])
        x_train, y_train = np.array(x_train), np.array(y_train)

        x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1], 1))

        # create and fit the LSTM network
        model = Sequential()
        model.add(LSTM(units=50, return_sequences=True, input_shape=(x_train.shape[1], 1)))
        model.add(LSTM(units=50))
        model.add(Dense(1))

        model.compile(loss='mean_squared_error', optimizer='adam')
        model.fit(x_train, y_train, epochs=1, batch_size=1, verbose=2)

        # predicting 246 values, using past 60 from the train data
        inputs = new_data[len(new_data) - len(valid) - 60:].values
        # inputs = inputs.reshape(-1, 1)
        inputs = scaler.transform(inputs)

        X_test = []
        for i in range(60, inputs.shape[0]):
            X_test.append(inputs[i - 60:i, 0])
        X_test = np.array(X_test)

        X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))
        closing_price = model.predict(X_test)
        # closing_price = scaler.inverse_transform(closing_price)

        new_valid = valid[:, 0:1]
        # rms = np.sqrt(np.mean(np.power((valid - closing_price), 2)))
        rms = np.sqrt(np.mean(np.power((new_valid - closing_price), 2)))
        Logger.info(rms)

        # for plotting
        train = new_data[:1500]
        valid = new_data[1500:]
        valid['Predictions'] = closing_price
        temp_train = train[0]
        plt.plot(temp_train)
        plt.plot(valid[[0, 'Predictions']])
        plt.show()

    def use_lstm(self, data):
        """
        开始运行LSTM算法
        :param data:
        :return:
        """
        Logger.info('开始运行LSTM算法')

        # 特征值个数，列数-1，因为最后一位为label
        n1 = len(data[0]) - 1
        # 行数
        n2 = len(data)
        Logger.info('特征值个数为【' + str(n1) + '】,行数为【' + str(n2) + '】')

        # 设置常量
        input_size = n1  # 输入神经元个数
        rnn_unit = 10    # LSTM单元(一层神经网络)中的中神经元的个数
        lstm_layers = 2  # LSTM单元个数
        output_size = 1  # 输出神经元个数(预测值)
        lr = 0.0006      # 学习率
        train_data_percentage = 0.9     # 训练数据百分比

        # 前90%数据作为训练集，后10%作为测试集
        train_end_index = math.floor(n2 * train_data_percentage)  # 向下取整
        Logger.info('训练数据的个数为【' + str(train_end_index) + '】')

        def get_train_data(batch_size=60, time_step=20, train_begin=0, train_end=train_end_index):
            """
            获取训练集
            time_step 时间步，batch_size 每一批次训练多少个样例
            :param batch_size:
            :param time_step:
            :param train_begin:
            :param train_end:
            :return:
            """
            batch_index = []
            data_train = data[train_begin:train_end]
            normalized_train_data = (data_train - np.mean(data_train, axis=0)) / np.std(data_train, axis=0)  # 标准化
            train_x, train_y = [], []  # 训练集
            for i in range(len(normalized_train_data) - time_step):
                if i % batch_size == 0:
                    # 开始位置
                    batch_index.append(i)
                    # 一次取time_step行数据
                # x存储输入维度&#xff08;不包括label&#xff09; :X(最后一个不取&#xff09;
                # 标准化(归一化&#xff09;
                x = normalized_train_data[i:i + time_step, :n1]
                # y存储label
                y = normalized_train_data[i:i + time_step, n1, np.newaxis]
                # np.newaxis分别是在行或列上增加维度
                train_x.append(x.tolist())
                train_y.append(y.tolist())
            # 结束位置
            batch_index.append((len(normalized_train_data) - time_step))
            print('batch_index', batch_index)
            # print('train_x', train_x)
            # print('train_y', train_y)
            return batch_index, train_x, train_y

        def get_test_data(time_step=20, test_begin=train_end_index+1):
            """
            获取测试集
            :param time_step:
            :param test_begin:
            :return:
            """
            data_test = data[test_begin:]
            mean = np.mean(data_test, axis=0)
            std = np.std(data_test, axis=0)  # 矩阵标准差
            # 标准化(归一化）
            normalized_test_data = (data_test - np.mean(data_test, axis=0)) / np.std(data_test, axis=0)
            # " // "表示整数除法。有size个sample
            test_size = (len(normalized_test_data) + time_step - 1) // time_step
            print('test_size$$$$$$$$$$$$$$', test_size)
            test_x, test_y = [], []
            for i in range(test_size - 1):
                x = normalized_test_data[i * time_step:(i+ 1) * time_step, :n1]
                y = normalized_test_data[i * time_step:(i + 1) * time_step, n1]
                test_x.append(x.tolist())
                test_y.extend(y)
            test_x.append((normalized_test_data[(i + 1) * time_step:, :n1]).tolist())
            test_y.extend((normalized_test_data[(i + 1) * time_step:, n1]).tolist())
            return mean, std, test_x, test_y

        # ——————————————————定义神经网络变量——————————————————
        # 输入层、输出层权重、偏置、dropout参数
        # 随机产生 w,b
        weights = {
            'in': tf.Variable(tf.random.normal([input_size, rnn_unit])),
            'out': tf.Variable(tf.random.normal([rnn_unit, 1]))
        }
        biases = {
            'in': tf.Variable(tf.constant(0.1, shape=[rnn_unit, ])),
            'out': tf.Variable(tf.constant(0.1, shape=[1, ]))
        }
        tf.compat.v1.disable_eager_execution()
        keep_prob = tf.compat.v1.placeholder(tf.float32, name='keep_prob')  # dropout 防止过拟合

        # ——————————————————定义神经网络——————————————————
        def lstmCell():
            """
            basicLstm单元
            tf.nn.rnn_cell.BasicLSTMCell(self, num_units, forget_bias=1.0,
            tate_is_tuple=True, activation=None, reuse=None, name=None)
            num_units:int类型，LSTM单元(一层神经网络)中的中神经元的个数，和前馈神经网络中隐含层神经元个数意思相同
            forget_bias:float类型，偏置增加了忘记门。从CudnnLSTM训练的检查点(checkpoin)恢复时，必须手动设置为0.0。
            state_is_tuple:如果为True，则接受和返回的状态是c_state和m_state的2-tuple；如果为False，则他们沿着列轴连接。后一种即将被弃用。
            （LSTM会保留两个state，也就是主线的state(c_state),和分线的state(m_state)，会包含在元组（tuple）里边
            state_is_tuple=True就是判定生成的是否为一个元组）
            初始化的 c 和 a 都是zero_state 也就是都为list[]的zero，这是参数state_is_tuple的情况下
            初始state,全部为0，慢慢的累加记忆
            activation:内部状态的激活函数。默认为tanh
            reuse:布尔类型，描述是否在现有范围中重用变量。如果不为True，并且现有范围已经具有给定变量，则会引发错误。
            name:String类型，层的名称。具有相同名称的层将共享权重，但为了避免错误，在这种情况下需要reuse=True.
            :return:
            """
            basicLstm = tf.compat.v1.nn.rnn_cell.BasicLSTMCell(rnn_unit, forget_bias=0.7, state_is_tuple=True)
            # dropout 未使用
            drop = tf.compat.v1.nn.rnn_cell.DropoutWrapper(basicLstm, output_keep_prob=keep_prob)
            return basicLstm

        def lstm(X):  # 参数：输入网络批次数目
            batch_size = tf.shape(X)[0]
            time_step = tf.shape(X)[1]
            w_in = weights['in']
            b_in = biases['in']

            # 忘记门（输入门）
            # 因为要进行矩阵乘法,所以reshape
            # 需要将tensor转成2维进行计算
            input = tf.reshape(X, [-1, input_size])
            input_rnn = tf.compat.v1.matmul(input, w_in) + b_in
            # 将tensor转成3维，计算后的结果作为忘记门的输入
            input_rnn = tf.reshape(input_rnn, [-1, time_step, rnn_unit])
            print('input_rnn', input_rnn)
            # 更新门
            # 构建多层的lstm
            cell = tf.compat.v1.nn.rnn_cell.MultiRNNCell([lstmCell() for i in range(lstm_layers)])
            init_state = cell.zero_state(batch_size, dtype=tf.float32)

            # 输出门
            w_out = weights['out']
            b_out = biases['out']
            # output_rnn是最后一层每个step的输出,final_states是每一层的最后那个step的输出
            output_rnn, final_states = tf.compat.v1.nn.dynamic_rnn(cell, input_rnn, initial_state=init_state, dtype=tf.float32)
            output = tf.reshape(output_rnn, [-1, rnn_unit])
            # 输出值，同时作为下一层输入门的输入
            pred = tf.matmul(output, w_out) + b_out
            return pred, final_states

        # ————————————————训练模型————————————————————
        def train_lstm(batch_size=60, time_step=20, train_begin=0, train_end=train_end_index):
            # 于是就有了tf.placeholder，
            # 我们每次可以将 一个minibatch传入到x = tf.placeholder(tf.float32,[None,32])上，
            # 下一次传入的x都替换掉上一次传入的，
            # 这样就对于所有传入的minibatch x就只会产生一个，
            # 不会产生其他多余的op&#xff0c;进而减少了graph的开销。

            X = tf.compat.v1.placeholder(tf.float32, shape=[None, time_step, input_size])
            Y = tf.compat.v1.placeholder(tf.float32, shape=[None, time_step, output_size])
            batch_index, train_x, train_y = get_train_data(batch_size, time_step, train_begin, train_end)
            # 用tf.variable_scope来定义重复利用,LSTM会经常用到
            with tf.compat.v1.variable_scope("sec_lstm"):
                pred, state_ = lstm(X) # pred输出值&#xff0c;state_是每一层的最后那个step的输出
            print('pred,state_', pred, state_)

            # 损失函数
            # [-1]——列表从后往前数第一列&#xff0c;即pred为预测值&#xff0c;Y为真实值(Label)
            # tf.reduce_mean 函数用于计算张量tensor沿着指定的数轴&#xff08;tensor的某一维度&#xff09;上的的平均值
            loss = tf.reduce_mean(tf.square(tf.reshape(pred, [-1]) - tf.reshape(Y, [-1])))
            # 误差loss反向传播——均方误差损失
            # 本质上是带有动量项的RMSprop&#xff0c;它利用梯度的一阶矩估计和二阶矩估计动态调整每个参数的学习率。
            # Adam的优点主要在于经过偏置校正后&#xff0c;每一次迭代学习率都有个确定范围&#xff0c;使得参数比较平稳.
            train_op = tf.compat.v1.train.AdamOptimizer(lr).minimize(loss)
            saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), max_to_keep=15)

            with tf.compat.v1.Session() as sess:
                # 初始化
                sess.run(tf.compat.v1.global_variables_initializer())
                theloss = []
                # 迭代次数
                for i in range(200):
                    for step in range(len(batch_index) - 1):
                        # sess.run(b, feed_dict = replace_dict)
                        state_, loss_ = sess.run([train_op, loss], feed_dict ={X: train_x[batch_index[step]:batch_index[step + 1]], Y: train_y[batch_index[step]:batch_index[step + 1]],
                                                         keep_prob: 0.5})
                        #  使用feed_dict完成矩阵乘法 处理多输入
                        #  feed_dict的作用是给使用placeholder创建出来的tensor赋值

                        #  [batch_index[step]: batch_index[step &#43; 1]]这个区间的X与Y
                        #  keep_prob的意思是&#xff1a;留下的神经元的概率&#xff0c;如果keep_prob为0的话&#xff0c; 就是让所有的神经元都失活。
                        print('Number of iterations', i, "loss", loss_)
                        theloss.append(loss_)
                print("model_save:", saver.save(sess, 'model_save2\\modle.ckpt'))
                print("The train has finished")
            return theloss

        theloss = train_lstm()

        # ————————————————预测模型————————————————————
        def prediction(time_step=20):
            X = tf.compat.v1.placeholder(tf.float32, shape=[None, time_step, input_size])
            mean, std, test_x, test_y = get_test_data(time_step)
            # 用tf.variable_scope来定义重复利用,LSTM会经常用到
            # with tf.compat.v1.variable_scope("sec_lstm", reuse=tf.AUTO_REUSE):
            with tf.compat.v1.variable_scope("sec_lstm", reuse=True):
                 pred, state_ = lstm(X)
            saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())
            with tf.compat.v1.Session() as sess:
                # 参数恢复（读取已存在模型）
                module_file = tf.train.latest_checkpoint('model_save2')
                saver.restore(sess, module_file)
                test_predict = []
                for step in range(len(test_x) - 1):
                    predict = sess.run(pred, feed_dict={X: [test_x[step]], keep_prob: 1})
                    predict = predict.reshape((-1))
                    test_predict.extend(predict)  # 把predict的内容添加到列表

                # 相对误差=（测量值-计算值）/计算值×100%
                test_y = np.array(test_y) * std[n1] + mean[n1]
                test_predict = np.array(test_predict) * std[n1] + mean[n1]
                acc = np.average(np.abs(test_predict - test_y[:len(test_predict)]) / test_y[:len(test_predict)])
                print("预测的相对误差", acc)

                print(theloss)
                plt.figure()
                plt.plot(list(range(len(theloss))), theloss, color ='b', )
                plt.xlabel('times', fontsize=14)
                plt.ylabel('loss valuet', fontsize=14)
                plt.title('loss-----blue', fontsize=10)
                plt.show()
                # 以折线图表示预测结果
                plt.figure()
                plt.plot(list(range(len(test_predict))), test_predict, color ='b', )
                plt.plot(list(range(len(test_y))), test_y, color ='r')
                plt.xlabel('time value/day', fontsize=14)
                plt.ylabel('close value/point', fontsize=14)
                plt.title('predict-----blue,real-----red', fontsize=10)
                plt.show()

        prediction()